#%%
### This script was written by Eric Deal, 06.22.2020
### Script should be run with python 3
import numpy as np, pandas as pd

# function to calculate settling velocity of sphere as per Dietrich et al,
def R1(logdx):
    # a, b, c, d, f = -3.76715, 1.92944, -0.09815, -0.00575, 0.00056
    a, b, c, d, f = -3.815642429922084, 1.9459392605305115, -0.09016149260089677, -0.008555578212166149, 0.000757318175522988
    # a, b, c, d, f = -16.835294914867273, 9.321922577683942, -1.6349454776208596, 0.13336298362521581, -0.00407146296540997
    return (a + b*logdx + c*logdx**2 + d*logdx**3 + f*logdx**4)

rof = 990.0
nu = 1e-6
g = 9.8

df_density = pd.read_csv('1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)

grains = df_density.index.values.astype('unicode')
dx_veqs = {grain: (df_density['total_grain_ro'][grain] - rof) * g * df_shape.dn[grain]**3 / (rof*nu*nu) for grain in grains}
ws_sphere = {grain: (10**R1(np.log10(dx_veqs[grain])) * (df_density['total_grain_ro'][grain] - rof)*g*nu/rof)**(1./3.) for grain in grains}

dx_veqs = {grain: (df_density['total_grain_ro'][grain] + df_density['dtotal_grain_ro'][grain] - rof) * g * df_shape.dn[grain]**3 / (rof*nu*nu) for grain in grains}
dws_sphere1 = {grain: (10**R1(np.log10(dx_veqs[grain])) * (df_density['total_grain_ro'][grain] + df_density['dtotal_grain_ro'][grain] - rof)*g*nu/rof)**(1./3.) for grain in grains}

dx_veqs = {grain: (df_density['total_grain_ro'][grain] - df_density['dtotal_grain_ro'][grain] - rof) * g * df_shape.dn[grain]**3 / (rof*nu*nu) for grain in grains}
dws_sphere2 = {grain: (10**R1(np.log10(dx_veqs[grain])) * (df_density['total_grain_ro'][grain] - df_density['dtotal_grain_ro'][grain] - rof)*g*nu/rof)**(1./3.) for grain in grains}

# dws_sphere = {grain: dws_sphere1[grain] - dws_sphere2[grain] for grain in grains}
dws_sphere = {grain: 0 for grain in grains}

df = pd.DataFrame({'ws_sphere': ws_sphere, 'dws_sphere': dws_sphere})

df_vol = pd.read_csv('1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)

df['CD_VES'] = np.nan
df['Re'] = np.nan
for grain in df.ws_sphere.index:
        df['CD_VES'][grain] = 4*(df_vol.total_grain_ro[grain]/990 - 1)*9.81*df_shape.dn[grain] / (3*df.ws_sphere[grain]**2) #* df_shape.CSF[grain]
        df['Re'][grain] = 1e6*df_shape.A[grain]

df.to_csv('5_var_VES_settling_velocity.csv')
